#!/bin/python
from __future__ import print_function

# derived from http://www.brool.com/index.php/using-google-authenticator-for-your-website

from base64 import b32encode, b32decode
from os import urandom
from time import time
from struct import pack, unpack
from hmac import HMAC
from hashlib import sha1
import logging

class TOTP:

	class ExceptionBadAuth:
		def __init__(self):
			pass

	_baseURL = "https://www.google.com/chart?chs=200x200&chld=M|0&cht=qr&chl=otpauth://totp/{}%3Fsecret%3D{}"
	_baseSimple = "otpauth://totp/{}?secret={}"
	ident = None
	secret = None
	url = None
	simple = None
	def __init__(self):
		logging.debug("TOTP init'ed")

	def generate(self, ident):
		self.ident = ident
		self.secret = b32encode(urandom(10))
		self.url = self._baseURL.format(self.ident, self.secret)
		self.simple = self._baseSimple.format(self.ident, self.secret)

	def auth(self, secretkey, attempt):
		logging.debug("Secretkey: ({}) attempt ({})".format(secretkey, attempt))
		tm = int(time() / 30)
		secretkey = b32decode(secretkey)

		# try 30 seconds behind and ahead as well
		for ix in [-1, 0, 1]:
			# convert timestamp to raw bytes
			b = pack(">q", tm + ix)
			
			# generate HMAC-SHA1 from timestamp based on secret key
			hm = HMAC(secretkey, b, sha1).digest()
			
			# extract 4 bytes from digest based on LSB
			offset = ord(hm[-1]) & 0x0F
			truncatedHash = hm[offset:offset+4]

			# get the code from it
			code = unpack(">L", truncatedHash)[0]
			code &= 0x7FFFFFFF;
			code %= 1000000;

			if ("%06d" % code) == str(attempt):
				logging.debug("Auth success")
				return True
		logging.debug("Auth fail")
		raise self.ExceptionBadAuth
		
if __name__ == "__main__":
	totp = TOTP()
	totp.generate(raw_input("Provide an identifier for this secret key: "))
	print("Secret Key: {}".format(totp.secret))
	print("URL: {}".format(totp.url))
	print("Simple: {}".format(totp.simple))
